/*
 * cube2video: convert cubes to videos (lossy or lossless).
 * GOU Lingfeng <goulf@ion.ac.cn>
 *
 * Only 8-bit and 16-bit (converted to 12-bit) data are supported.
 *
 */

#include <getopt.h>
#include <iostream>
#include <memory>
#include <sys/stat.h>
#include <string.h>
#include <libgen.h>

#include "public/utils.h"
#include "image.reader.h"
#include "config-fnt.h"

#include <vpx/vpx_encoder.h>
#include <vpx/vp8cx.h>
#include <mkvwriter.hpp>

struct ConvParams {
	uint64_t xleft, yleft, zleft;
	uint64_t xdelta, ydelta, zdelta;
	uint64_t xright, yright, zright;

	int quality, quality12; // [0 63]
	int preset, preset12; // [-8 8]
	bool use_msb;

	ConvParams():
		xleft{0}, yleft{0}, zleft{0},
		xdelta{0}, ydelta{0}, zdelta{0},
		xright{0}, yright{0}, zright{0},
		quality{24}, quality12{3},
		preset{-5}, preset12{-5},
		use_msb{false}
	{ }
	bool multipleCubes() {
		return xleft+xdelta<xright || yleft+ydelta<yright || zleft+zdelta<zright;
	}
};

/* Options */
static const char* opts=":x:y:z:mq:p:s:j::hv";
static const struct option opts_long[]=
{
	{"xrange", 1, nullptr, 'x'},
	{"yrange", 1, nullptr, 'y'},
	{"zrange", 1, nullptr, 'z'},
	{"msb", 0, nullptr, 'm'},
	{"quality", 1, nullptr, 'q'},
	{"preset", 1, nullptr, 'p'},
	{"stat", 1, nullptr, 's'},
	{"jobs", 2, nullptr, 'j'},
	{"help", 0, nullptr, 'h'},
	{"version", 0, nullptr, 'v'},
	{nullptr, 0, nullptr, 0}
};

/* Print version information and exit */
void version(const char* name) {
	std::cout<<name<<" (" PACKAGE_NAME_LONG ") " PACKAGE_VERSION "\n";
	std::cout<<"Copyright (C) 2015-2017 " PACKAGE_ORG "\n";
	std::cout<<"License GPLv3+: GNU GPL version 3 or later <http://gnu.org/licenses/gpl.html>\n\n";
	std::cout<<"Written by: " PACKAGE_AUTHORS "\n";
	exit(0);
}

/* Print usage information and exit */
void usage(const char* name, int ecode) {
	if(ecode) {
		std::cerr<<"Try '"<<name<<" --help' for more information.\n";
		exit(ecode);
	}
	std::cout<<"Usage: "<<name<<" [OPTION]... INPUT OUTPUT\n";
	std::cout<<
		"Convert cubes to videos for better compression.\n\n"
		"Options:\n"
		"  -h, --help       [OPTIONAL] Print this help, and do nothing else.\n"
		"  -v, --version    [OPTIONAL] Print version information and exit.\n\n"
		"Report bugs to <" PACKAGE_BUGREPORT ">\n"
		PACKAGE_NAME_LONG " home page: <" PACKAGE_URL ">\n";
	exit(ecode);
}

void createDirectory(const std::string& fulldir) {
	std::string tmpdir;
	for(size_t i=0; i<fulldir.size(); i++) {
		tmpdir.push_back(fulldir[i]);
		bool e=(i+1)==fulldir.size();
		bool s=!e && fulldir[i+1]=='/';
		if(e||s) {
			struct stat buf;
			int r=stat(tmpdir.c_str(), &buf);
			if(r==-1 || !S_ISDIR(buf.st_mode)) {
#ifdef BUILD_MINGW
				r=mkdir(tmpdir.c_str());
#else
				r=mkdir(tmpdir.c_str(), 0777);
#endif
				if(r==-1) {
					throwError("Cannot create dir '", tmpdir, "': ", strerror(errno), ".\n");
				}
			}

		}
	}
}

CubeType decode(std::vector<char>& buf, const char* input_file, int64_t* pw, int64_t* ph, int64_t* pd) {
	std::unique_ptr<ImageReader> imageReader{ImageReader::create(input_file)};
	if(!imageReader)
		throwError("Cannot read file: ", input_file);

	CubeType type;
	if(!imageReader->getType(&type))
		throwError("Cannot get image type: ", input_file);

	switch(type) {
		case CubeType::U8:
		case CubeType::U16:
			break;
		case CubeType::U32:
		case CubeType::I8:
		case CubeType::I16:
		case CubeType::I32:
		case CubeType::F32:
			throwError("Only 8/16 bit unsigned images supported.");
		default:
			throwError("Voxel type not supported");
	}

	int64_t w, h, d;
	if(!imageReader->getSizes(&w, &h, &d))
		throwError("Cannot get image sizes: ", input_file);

	auto bpv=bytes_per_voxel(type);
	try {
		if(buf.size()<uint64_t(bpv*w*h*d))
			buf.resize(bytes_per_voxel(type)*w*h*d);
	} catch(...) {
		throwError("Failed to alloc cube memory");
	}

	if(!imageReader->readData(buf.data(), 0, 0, 0, bpv*w, bpv*w*h))
		throwError("Failed to read image data: ", input_file);

	*pw=w;
	*ph=h;
	*pd=d;
	return type;
}


///////////////////////////////////////////////////////
void set_codec_pars(vpx_codec_ctx_t* ctx, int preset, int quality) {
	vpx_codec_err_t err;
	err=vpx_codec_control(ctx, VP8E_SET_CPUUSED, preset);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to set codec param: ", vpx_codec_err_to_string(err));
	//err=vpx_codec_control(ctx, VP8E_SET_TUNING, ?);
	//if(err!=VPX_CODEC_OK)
	//throwError("Failed to set codec param: ", vpx_codec_err_to_string(err));
	err=vpx_codec_control(ctx, VP8E_SET_CQ_LEVEL, quality);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to set codec param: ", vpx_codec_err_to_string(err));
	err=vpx_codec_control(ctx, VP9E_SET_FRAME_PARALLEL_DECODING, 1);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to set codec param: ", vpx_codec_err_to_string(err));
	//err=vpx_codec_control(ctx, VP9E_SET_ROW_MT, 1); // ???
	//if(err!=VPX_CODEC_OK)
	//throwError("Failed to set codec param: ", vpx_codec_err_to_string(err));
	//err=vpx_codec_control(ctx, VP9E_SET_LOSSLESS, 1);
	//if(err!=VPX_CODEC_OK)
		//throwError("Failed to set codec param: ", vpx_codec_err_to_string(err));
}

template<typename T> unsigned char* prepare_uv_plane(std::vector<char>& buf, int64_t width);
template<> unsigned char* prepare_uv_plane<uint8_t>(std::vector<char>& buf, int64_t width) {
	auto prev_size=buf.size();
	size_t size=(width+1)/2;
	if(prev_size<size)
		buf.resize(size, 0x80);
	return reinterpret_cast<unsigned char*>(&buf[0]);
}
template<> unsigned char* prepare_uv_plane<uint16_t>(std::vector<char>& buf, int64_t width) {
	auto prev_size=buf.size();
	size_t size=(width+1)/2*2;
	if(prev_size<size)
		buf.resize(size, 0x00);
	for(size_t i=prev_size+1; i<size; i+=2)
		buf[i]=0x08;
	return reinterpret_cast<unsigned char*>(&buf[0]);
}

template<typename T>
void encode(const char* output_fn, const char* stat_fn, unsigned char* cube, std::vector<char>& zerobuf, int64_t width, int64_t height, int64_t depth, const ConvParams& pars, int ncpu) {
	vpx_codec_err_t err;
	std::vector<char> buf_stats{};
	std::vector<char> buf_mb_stats{};
	vpx_image_t img;
	vpx_codec_enc_cfg_t cfg;
	vpx_codec_ctx_t ctx;

	vpx_codec_flags_t codec_flags=std::is_same<T, uint16_t>::value?VPX_CODEC_USE_HIGHBITDEPTH:0;
	auto preset=std::is_same<T, uint16_t>::value?pars.preset12:pars.preset;
	auto quality=std::is_same<T, uint16_t>::value?pars.quality12:pars.quality;

	vpx_codec_iface_t* iface=vpx_codec_vp9_cx();
	if(!iface)
		throwError("Failed to get vp9 interface.");

	memset(&img, 0, sizeof(img));
	img.img_data=cube;
	img.fmt=std::is_same<T, uint16_t>::value?VPX_IMG_FMT_I42016:VPX_IMG_FMT_I420;
	img.bit_depth=sizeof(T)*8;
	img.w=img.d_w=img.r_w=width;
	img.h=img.d_h=img.r_h=height;
	img.x_chroma_shift=1;
	img.y_chroma_shift=1;
	img.bps=std::is_same<T, uint16_t>::value?24:12;
	auto ystride=width*sizeof(T);
	img.stride[VPX_PLANE_Y]=ystride;
	img.stride[VPX_PLANE_ALPHA]=img.stride[VPX_PLANE_U]=img.stride[VPX_PLANE_V]=0;
	img.planes[VPX_PLANE_U]=img.planes[VPX_PLANE_V]=prepare_uv_plane<T>(zerobuf, width);
	//img.cs=VPX_CS_UNKNOWN;
	//img.range=VPX_CR_FULL_RANGE;

	err=vpx_codec_enc_config_default(iface, &cfg, 0);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to get config: ", vpx_codec_err_to_string(err));
	cfg.g_usage=VPX_CQ;
	cfg.g_threads=ncpu;
	cfg.g_profile=std::is_same<T, uint16_t>::value?2:0;
	cfg.g_w=width;
	cfg.g_h=height;
	cfg.g_bit_depth=std::is_same<T, uint16_t>::value?VPX_BITS_12:VPX_BITS_8;
	cfg.g_input_bit_depth=std::is_same<T, uint16_t>::value?12:8;
	cfg.g_timebase.num=50;
	cfg.g_timebase.den=1000;
	//cfg.g_lag_in_frames=;
	cfg.rc_dropframe_thresh=0;
	cfg.rc_resize_allowed=0;
	cfg.rc_end_usage=VPX_Q;
	cfg.kf_mode=VPX_KF_AUTO;
	cfg.kf_min_dist=0;
	cfg.kf_max_dist=0;

//#define VPX_TWO_PASS
#ifdef VPX_TWO_PASS
	cfg.g_pass=VPX_RC_FIRST_PASS;
	err=vpx_codec_enc_init(&ctx, iface, &cfg, codec_flags);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to init encoder: ", vpx_codec_err_to_string(err));
	set_codec_pars(&ctx, preset, quality);

	bool got_pkt=false;
	for(int64_t z=0; z<=depth || got_pkt; z++) {
		if(z<depth) {
			img.planes[VPX_PLANE_Y]=cube+ystride*height*z;
		}
		err=vpx_codec_encode(&ctx, z<depth?&img:nullptr, z, 1, VPX_EFLAG_FORCE_KF, 10);
		if(err!=VPX_CODEC_OK)
			throwError("Failed to encode frame: ", vpx_codec_err_to_string(err));

		got_pkt=false;
		vpx_codec_iter_t iter=nullptr;
		while(auto pkt=vpx_codec_get_cx_data(&ctx, &iter)) {
			const char* ptr;
			switch(pkt->kind) {
				case VPX_CODEC_STATS_PKT:
					ptr=static_cast<char*>(pkt->data.twopass_stats.buf);
					buf_stats.reserve(buf_stats.size()+pkt->data.twopass_stats.sz);
					buf_stats.insert(buf_stats.end(), ptr, ptr+pkt->data.twopass_stats.sz);
					got_pkt=true;
					break;
				case VPX_CODEC_FPMB_STATS_PKT:
					ptr=static_cast<char*>(pkt->data.firstpass_mb_stats.buf);
					buf_mb_stats.reserve(buf_mb_stats.size()+pkt->data.firstpass_mb_stats.sz);
					buf_mb_stats.insert(buf_mb_stats.end(), ptr, ptr+pkt->data.firstpass_mb_stats.sz);
					break;
				default:
					break;
			}
		}
	}
	err=vpx_codec_destroy(&ctx);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to destroy context: ", vpx_codec_err_to_string(err));
#endif

	FILE* webm_file=fopen(output_fn, "wb");
	if(!webm_file)
		throwError("Failed to open output file.");
	mkvmuxer::MkvWriter writer{webm_file};
	mkvmuxer::Segment segment{};
	segment.Init(&writer);
	segment.set_mode(mkvmuxer::Segment::kFile);
	segment.OutputCues(true);
	mkvmuxer::SegmentInfo* seginfo=segment.GetSegmentInfo();
	seginfo->set_timecode_scale(1000000);
	std::string verstr{PACKAGE_NAME_LONG " (libvpx-"};
	(verstr+=vpx_codec_version_str())+=")";
	seginfo->set_writing_app(verstr.c_str());
	auto track_id=segment.AddVideoTrack(width, height, 1);
	mkvmuxer::VideoTrack* track=static_cast<mkvmuxer::VideoTrack*>(segment.GetTrackByNumber(track_id));
	track->SetStereoMode(mkvmuxer::VideoTrack::kMono);
	track->set_codec_id("V_VP9");
	track->set_display_width(width);
	track->set_display_height(height);

#ifdef VPX_TWO_PASS
	cfg.g_pass=VPX_RC_LAST_PASS;
	cfg.rc_twopass_stats_in=vpx_fixed_buf_t{&buf_stats[0], buf_stats.size()};
	cfg.rc_firstpass_mb_stats_in=vpx_fixed_buf_t{&buf_mb_stats[0], buf_mb_stats.size()};
#else
	cfg.g_pass=VPX_RC_ONE_PASS;
#endif
	err=vpx_codec_enc_init(&ctx, iface, &cfg, codec_flags);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to init encoder: ", vpx_codec_err_to_string(err));
	set_codec_pars(&ctx, preset, quality);

	bool get_pkt=false;
	for(int64_t z=0; z<=depth || get_pkt; z++) {
		if(z<depth) {
			img.planes[VPX_PLANE_Y]=cube+ystride*height*z;
		}
		err=vpx_codec_encode(&ctx, z<depth?&img:nullptr, z, 1, VPX_EFLAG_FORCE_KF, 1);
		if(err!=VPX_CODEC_OK)
			throwError("Failed to encode frame: ", vpx_codec_err_to_string(err));

		get_pkt=false;
		vpx_codec_iter_t iter=nullptr;
		while(auto pkt=vpx_codec_get_cx_data(&ctx, &iter)) {
			int64_t pts_ns;
			switch(pkt->kind) {
				case VPX_CODEC_CX_FRAME_PKT:
					pts_ns=pkt->data.frame.pts*1000000000ll*cfg.g_timebase.num/cfg.g_timebase.den;
					segment.AddFrame(static_cast<uint8_t*>(pkt->data.frame.buf), pkt->data.frame.sz, 1, pts_ns, pkt->data.frame.flags&VPX_FRAME_IS_KEY);
					get_pkt=true;
					break;
				default:
					break;
			}
		}
	}
	err=vpx_codec_destroy(&ctx);
	if(err!=VPX_CODEC_OK)
		throwError("Failed to destroy context: ", vpx_codec_err_to_string(err));

	segment.Finalize();
	if(fclose(webm_file)!=0)
		throwError("Failed to close file.");

	if(stat_fn) {
		std::vector<char> buf(width*height*depth);
		int64_t w, h, d;
		decode(buf, output_fn, &w, &h, &d);
		std::vector<int64_t> hist(1<<(sizeof(T)*8), 0);
		for(int64_t i=0; i<width*height*depth; i++) {
			T diff=*(reinterpret_cast<T*>(&buf[0])+i)-*(reinterpret_cast<T*>(&cube[0])+i);
			hist[diff]++;
		}
		FILE* hf=fopen(stat_fn, "wb");
		for(size_t i=0; i<hist.size(); i++) {
			if(hist[i]) {
				fprintf(hf, "%zd %" PRId64 "\n", i, hist[i]);
			}
		}
		fclose(hf);
	}
}

void convert(const ConvParams& pars, const char* input_fn, const char* output_fn, const char* stat_fn, std::vector<char>& buf, std::vector<char>& zerobuf, int ncpu) {
	int64_t w, h, d;
	auto type=decode(buf, input_fn, &w, &h, &d);
	switch(type) {
		case CubeType::U8:
			encode<uint8_t>(output_fn, stat_fn, reinterpret_cast<unsigned char*>(buf.data()), zerobuf, w, h, d, pars, ncpu);
			break;
		case CubeType::U16:
			if(pars.use_msb) {
					auto ptr=reinterpret_cast<uint16_t*>(buf.data());
					auto ptre=reinterpret_cast<uint16_t*>(buf.data()+buf.size());
					while(ptr<ptre) {
				 *ptr=*ptr>>4;
				 //auto v=(*ptr)<<4;
				 //*ptr=(v>>8)|(v<<8);
				 ptr++;
				 }
			}
			encode<uint16_t>(output_fn, stat_fn, reinterpret_cast<unsigned char*>(buf.data()), zerobuf, w, h, d, pars, ncpu);
			break;
		case CubeType::U32:
		case CubeType::I8:
		case CubeType::I16:
		case CubeType::I32:
		case CubeType::F32:
			throwError("Only 8/16 bit unsigned images supported.");
		default:
			throwError("Voxel type not supported");
	}
}

void batchConvert(const ConvParams& pars, const char* input_pat, const char* output_pat, const char* stat_pat, int ncpu) {
	auto nx=(pars.xright-pars.xleft+pars.xdelta-1)/pars.xdelta;
	auto ny=(pars.yright-pars.yleft+pars.ydelta-1)/pars.ydelta;
	auto nz=(pars.zright-pars.zleft+pars.zdelta-1)/pars.zdelta;
	auto nn=nz*ny*nx;
#pragma omp parallel
	{
		std::vector<char> buf;
		std::vector<char> zerobuf;
#pragma omp for ordered schedule(static,1)
		for(uint64_t ni=0; ni<nn; ni++) {
			auto res=ni;
			auto x=(res%nx)*pars.xdelta+pars.xleft;
			res/=nx;
			auto y=(res%ny)*pars.ydelta+pars.yleft;
			res/=ny;
			auto z=(res%nz)*pars.zdelta+pars.zleft;
			res/=nz;

			auto input_fn=patternSubst(input_pat, x, y, z);
			auto output_fn=patternSubst(output_pat, x, y, z);
			auto output_path=output_fn;
			dirname(&output_path[0]);
			std::string stat_fn{};
			std::string stat_path{};
			if(stat_pat) {
				stat_fn=patternSubst(stat_pat, x, y, z);
				stat_path=stat_fn;
				dirname(&stat_path[0]);
			}
#pragma omp critical
			{
				std::cout<<ni<<'/'<<nn<<' '<<x<<':'<<y<<':'<<z<<'\n';
				createDirectory(output_path);
				if(stat_pat)
					createDirectory(stat_path);
			}

			convert(pars, input_fn.c_str(), output_fn.c_str(), stat_pat?stat_fn.c_str():nullptr, buf, zerobuf, 1);
		}
	}
}
int main(int argc, char* argv[]) {
	const char* input_pat=nullptr;
	const char* output_pat=nullptr;
	const char* stat_pat=nullptr;

	ConvParams pars{};
	int ncpu=1;

	int opt;
	int prev_optind=optind;
	while((opt=getopt_long(argc, argv, opts, opts_long, nullptr))!=-1) {
		try {
			switch(opt) {
				case '?':
					throwError("unrecognized option '", argv[prev_optind], "'\n");
				case ':':
					throwError("option '", argv[prev_optind], "' requires an argument\n");
				case 'h':
					usage(argv[0], 0);
				case 'v':
					version(argv[0]);
				case 'x':
					if(!parseTuple(optarg, &pars.xleft, &pars.xdelta, &pars.xright))
						throwError("invalid argument for '", argv[prev_optind], "'\nPlease give 3 integers separated by ':'.\n");
					if(pars.xleft>=pars.xright)
						throwError("invalid argument for '", argv[prev_optind], "'\nPlease give a valid range.\n");
					if(pars.xdelta==0)
						throwError("invalid argument for '", argv[prev_optind], "'\nSetep size must be positive.\n");
					break;
				case 'y':
					if(!parseTuple(optarg, &pars.yleft, &pars.ydelta, &pars.yright))
						throwError("invalid argument for '", argv[prev_optind], "'\nPlease give 3 integers separated by ':'.\n");
					if(pars.yleft>=pars.yright)
						throwError("invalid argument for '", argv[prev_optind], "'\nPlease give a valid range.\n");
					if(pars.ydelta==0)
						throwError("invalid argument for '", argv[prev_optind], "'\nSetep size must be positive.\n");
					break;
				case 'z':
					if(!parseTuple(optarg, &pars.zleft, &pars.zdelta, &pars.zright))
						throwError("invalid argument for '", argv[prev_optind], "'\nPlease give 3 integers separated by ':'.\n");
					if(pars.zleft>=pars.zright)
						throwError("invalid argument for '", argv[prev_optind], "'\nPlease give a valid range.\n");
					if(pars.zdelta==0)
						throwError("invalid argument for '", argv[prev_optind], "'\nSetep size must be positive.\n");
					break;
				case 'm':
					pars.use_msb=true;
					break;
				case 'q':
					if(!parseTuple<int>(optarg, &pars.quality))
						throwError("invalid argument for '", argv[prev_optind], "'.\nPlease give an integer.\n");
					if(pars.quality<0 || pars.quality>=64)
						throwError("invalid argument for '", argv[prev_optind], "'.\nThe range for quality is [0, 64).\n");
					pars.quality12=pars.quality;
					break;
				case 'p':
					if(!parseTuple<int>(optarg, &pars.preset))
						throwError("invalid argument for '", argv[prev_optind], "'.\nPlease give an integer.\n");
					if(pars.preset<-8 || pars.preset>8)
						throwError("invalid argument for '", argv[prev_optind], "'.\nThe range for preset is [-8, 8].\n");
					pars.preset12=pars.preset;
					break;
				case 's':
					stat_pat=optarg;
					break;
				case 'j':
					if(optarg) {
						if(!parseTuple<int>(optarg, &ncpu))
							throwError("invalid argument for '", argv[prev_optind], "'.\nPlease give an integer.\n");
						if(ncpu<=0)
							throwError("invalid argument for '", argv[prev_optind], "'.\nNeed a positive integer.\n");
					} else {
						ncpu=getNumberOfCpus();
						if(ncpu<1)
							ncpu=1;
					}
					break;
				default:
					throwError("Unexpected error.\n");
			}
		} catch(std::exception& e) {
			std::cerr<<argv[0]<<": "<<e.what();
			usage(argv[0], -1);
		}
		prev_optind=optind;
	}

	try {
		if(optind+2<argc)
			throwError("To many arguments.");
		if(optind+2>argc)
			throwError("Please specify input and output paths.");
		input_pat=argv[optind+0];
		output_pat=argv[optind+1];
	} catch(std::exception& e) {
		std::cerr<<argv[0]<<": "<<e.what()<<'\n';
		usage(argv[0], -1);
	}

	try {
		if(pars.multipleCubes()) {
			batchConvert(pars, input_pat, output_pat, stat_pat, ncpu);
		} else {
			std::vector<char> buf;
			std::vector<char> zerobuf;
			convert(pars, input_pat, output_pat, stat_pat, buf, zerobuf, ncpu);
		}
	} catch(std::exception& e) {
		std::cerr<<argv[0]<<": Conversion error: "<<e.what()<<'\n';
		usage(argv[0], -1);
	}

	return 0;
}

